﻿#ifndef __MAL_DOC__
#define __MAL_DOC__

#include "../comm/acsm.h"
#include "../comm/olefile.h"
#include "../../include/bitmap.h"
#include "../arch/pe/bptr.h"
#include "../arch/pe/hde32.h"

namespace xm {

	#define	BB_STR_PAIR(x)		x, sizeof(x)-1

	struct BBLOOKUP 
	{
		LPCVOID			pb;
		SIZE_T			cb;
	};

	static BBLOOKUP gShellCodeCookies[]=
	{
		BB_STR_PAIR("\x64\xA1"),
		BB_STR_PAIR("\x64\x8B"),
		BB_STR_PAIR("\xE8"),
		BB_STR_PAIR("\xD9\x74\x24\xF4"),
	};
	static BBLOOKUP gEncodeLookups[]=
	{
		BB_STR_PAIR("This program cannot be run in DOS mode"),
		BB_STR_PAIR(".text\x00\x00\x00"),
		BB_STR_PAIR(".data\x00\x00\x00"),
		BB_STR_PAIR(".rdata\x00\x00"),
		BB_STR_PAIR(".reloc\x00\x00"),
		BB_STR_PAIR(".rsrc\x00\x00\x00"),
		BB_STR_PAIR(".UPX0\x00\x00\x00"),
		BB_STR_PAIR("GetModuleFileName\x00"),
		BB_STR_PAIR("URLDownloadFile\x00"),
		BB_STR_PAIR("ShellExecute\x00"),
		BB_STR_PAIR("WinExec\x00"),
		BB_STR_PAIR("VirtualAlloc\x00"),
		BB_STR_PAIR("WriteProcessMemory\x00"),
		BB_STR_PAIR("CreateRemoteThread\x00"),
		BB_STR_PAIR("\x56\x8B\x75\x3C\x8B\x74\x2E\x78\x03\xF5\x56\x8B\x76\x20"),
		/*
		seg000:00006E00 8B 76 0C                                mov     esi, [esi+0Ch]
		seg000:00006E03 8B 76 1C                                mov     esi, [esi+1Ch]
		*/
		BB_STR_PAIR("\x64\x8B\x71\x30\x8B\x76\x0C\x8B\x76\x1C"),
		/*
		seg000:000070A8 8B 73 3C                                mov     esi, [ebx+3Ch]
		seg000:000070AB 8B 74 1E 78                             mov     esi, [esi+ebx+78h
		*/
		BB_STR_PAIR("\x8B\x73\x3C\x8B\x74\x1E\x78"),
		//BB_STR_PAIR("\x6A\x00\x6A\x00\x6A\x00"),
		//BB_STR_PAIR("\xD0\xCF\x11\xE0\xA1\xB1"),
		BB_STR_PAIR("\x56\x8B\x73\x3C\x8B\x74\x1E\x78\x01\xDE\x56\x8B\x76\x20"),
		BB_STR_PAIR("\x80\x38\xE8\x74\x0F\x80\x38\xE9\x74\x0A\x80\x38\xCC\x74\x05\x80\x38\xEB\x75\x11"),
		BB_STR_PAIR("\x64\x8B\x40\x30\x8B\x40\x0C\x8B\x40\x1C"),
		BB_STR_PAIR("\x68\x6C\x6C\x00\x00\x68\x6F\x6E\x2E\x64\x68\x75\x72\x6C\x6D"),
	};

	class MalDocACMs
	{
	public:
		mpm::acsm2_t acmXOR, acmSUB, acmRAW;
	protected:
		static SIZE_T GenEncPatternXOR( LPVOID out, LPCVOID p, SIZE_T c )
		{
			UINT8 * bp = (UINT8*)out;
			CONST UINT8 * bpi = (UINT8*)p;
			for( SIZE_T i = 1; i < c; ++ i )
				bp[i-1] = bpi[i-1] ^ bpi[i];
			return c-1;
		}
		static SIZE_T GenEncPatternSUB( LPVOID out, LPCVOID p, SIZE_T c )
		{
			UINT8 * bp = (UINT8*)out;
			CONST UINT8 * bpi = (UINT8*)p;
			for( SIZE_T i = 1; i < c; ++ i )
				bp[i-1] = bpi[i-1] - bpi[i];
			return c-1;
		}
	public:
		HRESULT Build()
		{
			// shellcode
			{
				mpm::acsm_t<> acm;
				for( SIZE_T i = 0; i < sizeof(gShellCodeCookies)/sizeof(gShellCodeCookies[0]); ++ i )
				{
					long lr = acm.insert( gShellCodeCookies[i].pb, gShellCodeCookies[i].cb, (void*)(i+1), false );
					if( lr < 0 ) return E_FAIL;
				}
				if( acm.complie() < 0 ) 
					return E_FAIL;
				if( acmRAW.convertFrom( acm.root() ) < 0 )
					return E_FAIL;
			}
			// XOR EncodeLookup
			UINT8 tmp[256] = {};
			{
				mpm::acsm_t<> acm;
				for( SIZE_T i = 0; i < sizeof(gEncodeLookups)/sizeof(gEncodeLookups[0]); ++ i )
				{
					SIZE_T c = GenEncPatternXOR( tmp, gEncodeLookups[i].pb, gEncodeLookups[i].cb );
					if( acm.insert( tmp, c, (void*)(i+1), false ) < 0 )
						return E_FAIL;
				}
				if( acm.complie() < 0 ) 
					return E_FAIL;
				if( acmXOR.convertFrom( acm.root() ) < 0 )
					return E_FAIL;
			}
			{
				mpm::acsm_t<> acm;
				for( SIZE_T i = 0; i < sizeof(gEncodeLookups)/sizeof(gEncodeLookups[0]); ++ i )
				{
					SIZE_T c = GenEncPatternSUB( tmp, gEncodeLookups[i].pb, gEncodeLookups[i].cb );
					if( acm.insert( tmp, c, (void*)(i+1), false ) < 0 )
						return E_FAIL;
				}
				if( acm.complie() < 0 ) 
					return E_FAIL;
				if( acmSUB.convertFrom( acm.root() ) < 0 )
					return E_FAIL;
			}
			return S_OK;
		}

	};

	static inline int memcmp_xor8( const void * o, const void * e, size_t c, int k )
	{
		const uint8_t _k = (const uint8_t)k;
		const uint8_t * _l = (const uint8_t *)o;
		const uint8_t * _r = (const uint8_t *)e;
		const uint8_t * _le = _l + c;
		for( ; _l < _le; ++ _l, ++ _r )
		{
			int diff = (*_l)^_k;
			diff -= *_r;
			if( diff ) return diff;
		}
		return 0;
	}
	static inline int memcmp_sub8( const void * o, const void * e, size_t c, int k )
	{
		const uint8_t _k = (const uint8_t)k;
		const uint8_t * _l = (const uint8_t *)o;
		const uint8_t * _r = (const uint8_t *)e;
		const uint8_t * _le = _l + c;
		for( ; _l < _le; ++ _l, ++_r )
		{
			int diff = (*_l) + _k;
			diff -= *_r;
			if( diff ) return diff;
		}
		return 0;
	}

	class MalDocDetectSession
	{
	protected:
		typedef std::vector<UINT8> BBuffer;
		blob_t&		_ScanData;
		SIZE_T		_Score;
		BBuffer		_XORDelta, _SubDelta;
		SIZE_T		_ScanBeginPos;
		UINT8		_TempBuffer[0x100];
		bit_map_t	_hitmap;
	public:

		MalDocDetectSession( blob_t& rBlob, SIZE_T aScanBeginPos ) : _ScanData(rBlob), _Score(0), _ScanBeginPos(aScanBeginPos)
		{

		}

	protected:

		size_t		FeatureCount()
		{
			return sizeof(gEncodeLookups)/sizeof(gEncodeLookups[0]);
		}

		static bool _on_enc_found_xor(const uint8_t * dat, size_t cb, size_t offset,void * ptnud,void * context )
		{
			return ((MalDocDetectSession*)context)->on_enc_found_xor( dat, cb, offset, ptnud );
		}
		static bool _on_enc_found_sub(const uint8_t * dat, size_t cb, size_t offset,void * ptnud,void * context )
		{
			return ((MalDocDetectSession*)context)->on_enc_found_sub( dat, cb, offset, ptnud );
		}
		static bool _on_hit_shellcode(const uint8_t * dat, size_t cb, size_t offset,void * ptnud,void * context )
		{
			return ((MalDocDetectSession*)context)->on_hit_shellcode( dat, cb, offset, ptnud );
		}
		bool on_hit_shellcode( CONST UINT8 * dat, SIZE_T cb, SIZE_T offset, LPVOID ptnud )
		{
			SIZE_T i = (SIZE_T)ptnud-1;
			if( 0 == i )
			{
				if( cb - offset >= 16 )
				{
					UINT32* fs_addr = (UINT32*)(dat + offset + 1);
					if( (fs_addr[0] == 0x00 || fs_addr[0] == 0x30) && fs_addr[1] != 0 )
					{
						_Score ++;
						//printf(" [%08X] mov eax, fs[00/30] found.\n", offset + _ScanBeginPos );
					}
				}
			}
			else if( 1 == i )
			{
				bptr_t bp( dat + offset - 1, dat + cb );
				hde32s ins;
				int il = hde32_disasm( bp, &ins );
				if( il < 3 ) return true;
				if( ins.disp.disp32 == 0x30 )
				{
					_Score ++;
					//printf(" [%08X] mov ???, fs[00/30] found.\n", offset + _ScanBeginPos );
				}
			}
			else if( 2 == i )
			{
				if( cb - offset >= 5 )
				{
					INT32 joff = *(INT32*)(dat + offset + 1);
					if( joff < -0x200 || joff > 0x200 ) return true;
					SIZE_T target = offset + joff + 5;
					if( target + 16 < cb )
					{
						if( (dat[target] & 0xF8) == 0x58 )
						{
							size_t z = 0;
							for( size_t i = 0; i < 10; ++ i )
							{
								if( !dat[target+i] )
									z ++;
							}
							if( z > 3 ) return true;
							//printf(" [%08X] call/pop ??? found.\n", offset + _ScanBeginPos );
							_Score ++;
						}
						if( (dat[target] == 0x8B) && ((dat[target+1] & 7) == 4) && (dat[target+2] == 0x24) )
						{
							//printf(" [%08X] call/mov ???, [esp] found.\n", offset + _ScanBeginPos );
							_Score ++;
						}
					}
				}
			}
			else if( 3 == i )
			{
				if( cb - offset >= 16 )
				{
					if( (dat[offset+1] & 0xF8) == 0x58 ) // pop
					{
						size_t z = 0;
						for( size_t i = 0; i < 10; ++ i )
						{
							if( !dat[offset+i] )
								z ++;
						}
						if( z > 3 ) return true;
						//printf(" [%08X] fnstenv byte ptr [esp-0Ch]/pop xxx found.\n", offset );
						_Score += 2;			
					}
				}
			}
			else
			{
				_Score ++;
			}
			return true;
		}
		bool on_enc_found_xor(CONST UINT8 * dat, SIZE_T cb, SIZE_T enc_offset, LPVOID ptnud )
		{
			SIZE_T iLookup = (SIZE_T)ptnud-1;
			if( _hitmap.testbit( iLookup ) ) return true;
			_hitmap.setbit( iLookup, 1 );
			CONST UINT8* bpPattern = (CONST UINT8*)gEncodeLookups[iLookup].pb;
			SIZE_T cbPattern = gEncodeLookups[iLookup].cb;
			CONST UINT8 * bpHitPtr = (CONST UINT8*)_ScanData.Data + enc_offset + 1 - cbPattern + 1;
			UINT8 bXOR = (*bpPattern) ^ (*bpHitPtr);
			if( memcmp_xor8( bpPattern, bpHitPtr, cbPattern, bXOR ) )
				return true;
			//printf("[%08X] found pattern %d:%s, XOR key = %02X\n", enc_offset + _ScanBeginPos + 1,  ptnud, bpPattern, bXOR );
			if( bXOR ) _Score ++;
			_Score ++;
			if( _Score > 2 ) return false;
			return true;
		}
		bool on_enc_found_sub(CONST UINT8 * dat, SIZE_T cb, SIZE_T enc_offset, LPVOID ptnud )
		{
			SIZE_T iLookup = (SIZE_T)ptnud-1;
			if( _hitmap.testbit( iLookup ) ) return true;
			_hitmap.setbit( iLookup, 1 );
			CONST UINT8* bpPattern = (CONST UINT8*)gEncodeLookups[iLookup].pb;
			SIZE_T cbPattern = gEncodeLookups[iLookup].cb;
			CONST UINT8 * bpHitPtr = (CONST UINT8*)_ScanData.Data + enc_offset + 1 - cbPattern + 1;
			UINT8 bSub = (*bpHitPtr) - (*bpPattern);
			if( memcmp_sub8( bpPattern, bpHitPtr, cbPattern, bSub ) )
				return true;
			//printf("[%08X] found pattern %d:%s, SUB key = %02X\n", enc_offset + _ScanBeginPos + 1,  ptnud, bpPattern, bSub );
			if( bSub ) _Score ++;
			_Score ++;
			if( _Score > 2 ) return false;
			return true;
		}
		inline CONST UINT8* ScanDataPtr()
		{
			return (CONST UINT8*)_ScanData.Data;
		}
		inline SIZE_T ScanDataSize()
		{
			return _ScanData.Length;
		}
		HRESULT GenLoopupBufferXOR()
		{
			CONST UINT8* bpScanData = ScanDataPtr();
			SIZE_T cbScanData = ScanDataSize();
			_XORDelta.reserve( cbScanData - 1 );
			for( SIZE_T i = 1; i < cbScanData; ++ i )
			{
				UINT8 tmp = bpScanData[i-1] ^ bpScanData[i];
				_XORDelta.push_back( tmp );
			}
			return S_OK;
		}
		HRESULT GenLoopupBufferSUB()
		{
			CONST UINT8* bpScanData = ScanDataPtr();
			SIZE_T cbScanData = ScanDataSize();
			_SubDelta.reserve( cbScanData - 1 );
			for( SIZE_T i = 1; i < cbScanData; ++ i )
			{
				UINT8 tmp = bpScanData[i-1] - bpScanData[i];
				_SubDelta.push_back( tmp );
			}
			return S_OK;
		}

	public:

		HRESULT Reset()
		{
			return _hitmap.resize( FeatureCount() ) ? S_OK : E_FAIL;
		}

		SIZE_T	GetScore()
		{
			return _Score;
		}

		HRESULT DetectInShellCode( mpm::acsm2_t& acm )
		{
			mpm::acsm2_t::node2* cnode = 0;
			long lr = acm.search( cnode, ScanDataPtr(), ScanDataSize(), &MalDocDetectSession::_on_hit_shellcode, this );
			if( lr < 0 ) return E_FAIL;
			return S_OK;
		}
		HRESULT DetectInXOR( mpm::acsm2_t& acm )
		{
			RFAILED( GenLoopupBufferXOR() );
			mpm::acsm2_t::node2* cnode = 0;
			long lr = acm.search( cnode, _XORDelta.data(), _XORDelta.size(), &MalDocDetectSession::_on_enc_found_xor, this );
			if( lr < 0 ) return E_FAIL;
			return S_OK;
		}
		HRESULT DetectInSUB( mpm::acsm2_t& acm )
		{
			RFAILED( GenLoopupBufferSUB() );
			mpm::acsm2_t::node2* cnode = 0;
			long lr = acm.search( cnode, _SubDelta.data(), _SubDelta.size(), &MalDocDetectSession::_on_enc_found_sub, this );
			if( lr < 0 ) return E_FAIL;
			return S_OK;
		}
	};

	class XdoBinExploit1 : public IKDetector, UnkBase
	{
	protected:
		MalDocACMs	_ACM;
	public:
		UNKNOWN_IMP1(IKDetector);
		HRESULT init_class( void * ,void * )
		{
			return _ACM.Build();
		}
		HRESULT Detect( blob_t & rFileData, SIZE_T cSkip = 0x80 )
		{
			if( rFileData.Length < cSkip ) return 0;

			blob_t ScanData;
			ScanData.Data = (LPCSTR)rFileData.Data + cSkip ;
			ScanData.Length = rFileData.Length - cSkip ;

			MalDocDetectSession thisScan( ScanData, cSkip );

			RFAILED( thisScan.Reset() );
			RFAILED( thisScan.DetectInShellCode( _ACM.acmRAW ) );
			if( thisScan.GetScore() > 1 ) return S_OK;

			RFAILED( thisScan.DetectInXOR( _ACM.acmXOR ) );
			if( thisScan.GetScore() > 1 ) return S_OK;

			RFAILED( thisScan.DetectInSUB( _ACM.acmSUB ) );
			if( thisScan.GetScore() > 1 ) return S_OK;
			return E_FAIL;
		}
		STDMETHOD(Detect)( IKDetectSession * sess, DetectResult * pdr )
		{
			optr<IKArchive> oArchive( sess->CurrentArchive() );
			if( !oArchive ) return E_INVALIDARG;
		
			FMTID uFormatId = 0;
			RFAILED( oArchive->GetFormat( &uFormatId ) );

			optr<IKBStream> oStream = sess->TargetStream();

			SIZE_T cbSkips = 0x80;

			if( uFormatId == FMT_DOC )
			{
				OleFileEx oOleFile;
				if( !oOleFile.OpenOleFile( oStream ) )
					return E_FAIL;

				BOOL bIsOfficeDocument = FALSE;
				DWORD DirNum = oOleFile.GetDirEntryNum();
				for( DWORD i = 0; i < DirNum; ++ i )
				{
					StorageDirEntry * lpEntry = oOleFile.GetDirEntry( i );
				
					CONST SIZE_T ch16MaxName = sizeof(lpEntry->bName)/sizeof(UINT16);
				
					WCHAR wName[ch16MaxName] = {};

					UINT16* p16Name = (UINT16*)lpEntry->bName;
				
					for( SIZE_T i16 = 0; i16 < ch16MaxName; ++ i16 ) 
						wName[i16] = p16Name[i16];

					if( !tcsicmp( wName, L"WordDocument" ) ||
						!tcsicmp( wName, L"WorkBook" ) ||
						!tcsicmp( wName, L"Contents" ) ||						
						!tcsicmp( wName, L"PowerPoint Document" ) )
					{
						bIsOfficeDocument = TRUE;
						break;
					}
				}
				if( !bIsOfficeDocument )
				{
					return E_FAIL;
				}
				else
				{
					cbSkips = 0x200;
				}
			}
			else
			{
			}

			SIZE_T mMax2Scan = 2*1024*1024;
			UTIL::sentry<UINT8*> buffer = new UINT8[mMax2Scan];

			SIZE_T eff = 0;
			RFAILED( oStream->Readp( cbSkips, buffer, mMax2Scan, &eff ) );

			blob_t bbFile;
			bbFile.Data = buffer;
			bbFile.Length = eff;

			RFAILED( Detect( bbFile ) ); 

			pdr->Level = DLEV_MALW_HEUR;
			return S_OK;
		}

		STDMETHOD(Treat)( IKDetectSession * sess, DetectResult * pdr )
		{
			return sess->RawStream()->Remove();
		}

		STDMETHOD(Reset)()
		{
			return S_OK;
		}

	};
};
#endif